#!/usr/bin/env python3
import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.executors import ExternalShutdownException
from rclpy.qos import qos_profile_sensor_data
from geometry_msgs.msg import PoseStamped, Pose, TransformStamped
from tf2_ros import Buffer, TransformListener, TransformBroadcaster
from tf2_ros.transformations import quaternion_matrix, quaternion_from_matrix
from scipy.spatial.transform import Rotation as R


# 定义旋转矩阵函数
def get_transform(tf_buffer, target_frame, source_frame):
    """
    获取指定两个坐标系之间的变换
    """
    try:
        transform = tf_buffer.lookup_transform(target_frame, source_frame, rclpy.time.Time())
        translation = np.array([transform.transform.translation.x,
                                transform.transform.translation.y,
                                transform.transform.translation.z])
        rotation = np.array([transform.transform.rotation.x,
                             transform.transform.rotation.y,
                             transform.transform.rotation.z,
                             transform.transform.rotation.w])
        R_matrix = quaternion_matrix(rotation)[:3, :3]
        return translation, R_matrix
    except Exception as e:
        print(f"Transform error: {e}")
    return None, None


def compute_grasp_position(P_object, R_object, shape, dimensions=None):
    """
    计算夹爪接触点和姿态
    Parameters:
    - P_object: 物体位置 (3x1 array)
    - R_object: 物体姿态 (3x3 rotation matrix)
    - shape: 物体形状 ('cylinder' or 'cuboid')
    - dimensions: 物体尺寸 (tuple), cylinder: (r, h), cuboid: (l, w, h)
    Returns:
    - P_grasp: 抓取点的位置
    - grasp_orientation: 抓取姿态 (3x3 rotation matrix)
    """
    if shape == 'cylinder':
        r, h = dimensions
        obj_x_dir = R_object[:, 0]
        obj_y_dir = R_object[:, 1]
        obj_z_dir = R_object[:, 2]

        tool_x_dir = -obj_z_dir
        tool_x_dir = tool_x_dir / np.linalg.norm(tool_x_dir)
        tool_y_candidate = np.array([0.0, -1.0, 0.0])

        dot_xy = np.abs(np.dot(tool_x_dir, tool_y_candidate))

        if dot_xy > 0.99:
            tool_y_candidate = np.array([0.0, 0.0, -1.0])

        v_perp = tool_y_candidate - np.dot(tool_y_candidate, tool_x_dir) * tool_x_dir
        norm_vp = np.linalg.norm(v_perp)

        if norm_vp < 1e-8:
            tool_y_candidate = np.array([0.0, 0.0, 1.0])
            v_perp = tool_y_candidate - np.dot(tool_y_candidate, tool_x_dir) * tool_x_dir
            norm_vp = np.linalg.norm(v_perp)

        tool_y_dir = v_perp / norm_vp
        tool_z_dir = np.cross(tool_x_dir, tool_y_dir)
        tool_z_dir = tool_z_dir / np.linalg.norm(tool_z_dir)

        grasp_orientation = np.column_stack((tool_x_dir, tool_y_dir, tool_z_dir))
        grasp_poses = [{
            "position": P_object + grasp_orientation @ np.array([0, r, 0.063]),
            "orientation": grasp_orientation
        }]

        br = TransformBroadcaster()
        t = TransformStamped()
        t.header.stamp = rclpy.time.Time().to_msg()
        t.header.frame_id = "base_link"
        t.child_frame_id = "target_tool_link"

        t.transform.translation.x = grasp_poses[0]["position"][0]
        t.transform.translation.y = grasp_poses[0]["position"][1]
        t.transform.translation.z = grasp_poses[0]["position"][2]

        R_object_orientation = np.eye(4)
        R_object_orientation[:3, :3] = grasp_poses[0]["orientation"]
        t.transform.rotation.x = quaternion_from_matrix(R_object_orientation)[0]
        t.transform.rotation.y = quaternion_from_matrix(R_object_orientation)[1]
        t.transform.rotation.z = quaternion_from_matrix(R_object_orientation)[2]
        t.transform.rotation.w = quaternion_from_matrix(R_object_orientation)[3]
        br.sendTransform(t)

    elif shape == 'cuboid':
        l, w, h = dimensions
        P_contact = np.array([l / 2, 0, 0])
        P_grasp = P_object + np.dot(R_object, P_contact)
        x_axis = np.array([1, 0, 0])
        z_axis = R_object[:, 2]
        y_axis = np.cross(z_axis, x_axis)
        y_axis /= np.linalg.norm(y_axis)
        z_axis = np.cross(x_axis, y_axis)
        grasp_orientation = np.column_stack((x_axis, y_axis, z_axis))
    else:
        raise ValueError("Unsupported shape: choose 'cylinder' or 'cuboid'")

    return grasp_poses[0]["position"], grasp_poses[0]["orientation"]


def apply_fixed_transform(P_object, R_object, transform_params):
    """
    应用固定的变换参数
    """
    qw = transform_params["qw"]
    qx = transform_params["qx"]
    qy = transform_params["qy"]
    qz = transform_params["qz"]
    x = transform_params["x"]
    y = transform_params["y"]
    z = transform_params["z"]

    fixed_rotation = np.array([qx, qy, qz, qw])
    R_fixed = quaternion_matrix(fixed_rotation)[:3, :3]
    P_fixed = np.array([x, y, z])

    P_object_transformed = np.dot(R_fixed, P_object) + P_fixed
    R_object_transformed = np.dot(R_fixed, R_object)

    return P_object_transformed, R_object_transformed


def process_pose_data(msg, grasp_publisher, tf_buffer):
    """
    处理从PoseStamped消息中接收到的pose数据并计算抓取位姿
    """
    P_object = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
    quaternion = [msg.pose.orientation.x, msg.pose.orientation.y, msg.pose.orientation.z, msg.pose.orientation.w]
    R_object = quaternion_matrix(quaternion)[:3, :3]

    shape = 'cylinder'
    dimensions = [0.070, 0.106]  # 直径和高度

    P_grasp, grasp_orientation = compute_grasp_position(P_object, R_object, shape, dimensions)

    grasp_orientation_homogeneous = np.eye(4)
    grasp_orientation_homogeneous[:3, :3] = grasp_orientation
    grasp_quaternion = quaternion_from_matrix(grasp_orientation_homogeneous)

    grasp_msg = Pose()
    grasp_msg.position.x = P_grasp[0]
    grasp_msg.position.y = P_grasp[1]
    grasp_msg.position.z = P_grasp[2]
    grasp_msg.orientation.x = grasp_quaternion[0]
    grasp_msg.orientation.y = grasp_quaternion[1]
    grasp_msg.orientation.z = grasp_quaternion[2]
    grasp_msg.orientation.w = grasp_quaternion[3]

    grasp_publisher.publish(grasp_msg)
    print(f"Grasp Pose Published: {grasp_msg}")


class GraspCalculationNode(Node):
    def __init__(self):
        super().__init__("grasp_calculation_node")
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer)
        self.grasp_publisher = self.create_publisher(Pose, "/right_grasp_pose", qos_profile_sensor_data)
        self.create_subscription(PoseStamped, "/obj_pose", lambda msg: process_pose_data(msg, self.grasp_publisher, self.tf_buffer), qos_profile_sensor_data)


def main(args=None):
    rclpy.init(args=args)
    try:
        node = GraspCalculationNode()
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    except ExternalShutdownException:
        pass
    finally:
        rclpy.try_shutdown()
        print("Program exited.")


if __name__ == "__main__":
    main()
